import numpy as np
import pandas as pd
from fontTools.unicodedata import block
from matplotlib import pyplot as plt
import Log

# 配置和初始化 logger
logger = Log.setup_logger()
trade_rate = 0.86 / 10000

gTop = 0
gBottom = 0
gR = -1000
gStep = []
gMoney = []


def fund(firstCode, secondCode, needBackTest=True, needDrawPic=True, days=4000):
    disList = []
    dateList = []
    dfSecond = None
    dfFirst = None
    returnDiffDay = 40
    log(f'fund {firstCode}-{secondCode} returnDiffDay({returnDiffDay})')
    dfSecond = pd.read_csv(f'data/{secondCode}.csv', index_col=0, parse_dates=True)
    dfSecond = dfSecond.fillna(method="bfill")
    # dfSecond = dfSecond[(dfSecond.index > '2023-01-01') & (dfSecond.index < '2023-12-31')]
    # dfSecond = dfSecond[(dfSecond.index < '2023-12-31')]
    dfFirst = pd.read_csv(f'data/{firstCode}.csv', index_col=0, parse_dates=True)
    dfFirst = dfFirst.fillna(method="bfill")
    # dfFirst = dfFirst[(dfFirst.index > '2023-01-01') & (dfFirst.index < '2023-12-31')]
    # dfFirst = dfFirst[(dfFirst.index < '2023-12-31')]

    closeListSecond = dfSecond['close_price']
    numSecond = closeListSecond.size
    # 科创200
    closeListFirst = dfFirst['close_price']
    if closeListFirst.index[-1] != closeListSecond.index[-1]:
        raise Exception('not same end date')
    numFirst = closeListFirst.size
    diffSize = min(numSecond, numFirst, days)
    log(f'diff size {diffSize}')
    closeListSecond = closeListSecond.tail(diffSize)
    closeListFirst = closeListFirst.tail(diffSize)
    if needDrawPic:
        drawBollLine(closeListFirst, closeListSecond)
    for i in range(0, diffSize - returnDiffDay + 1):
        dayBefore40 = closeListSecond[i]
        currentDay = closeListSecond[i + returnDiffDay - 1]
        yieldValueKC50 = (currentDay - dayBefore40) / dayBefore40
        dayBefore40 = closeListFirst[i]
        currentDay = closeListFirst[i + returnDiffDay - 1]
        yieldValueKC200 = (currentDay - dayBefore40) / dayBefore40
        disList.append(round((yieldValueKC200 - yieldValueKC50) * 100, 2))
        dateList.append(closeListSecond.index[i + returnDiffDay - 1].date())

    # 绘制折线图
    if needDrawPic:
        drawPic(firstCode, secondCode, disList, dateList)
    # backTest(8, -12, disList,dateList, closeListFirst, closeListSecond,returnDiffDay)
    if needBackTest:
        for i in range(1, int(max(disList)) + 1):
            for j in range(int(min(disList)), 0):
                backTest(i, j, disList, dateList, closeListFirst, closeListSecond, returnDiffDay)

        logger.info(
            f'{firstCode}-{secondCode} 最大收益率({round(gR * 100, 2)}%) 差值上沿({gTop}%) 差值下沿({gBottom}%) 年化收益率({calculate_annualized_return(int(gR * 100), diffSize - returnDiffDay)}%)')
        for i in range(0, len(gStep)):
            log(gStep[i])
    if needDrawPic:
        drawMoneyPic(firstCode, secondCode, gMoney, dateList)
        log(
            f"最大回撤({calculate_max_drawdown(gMoney)}) 第一个最大回撤({calculate_max_drawdown(closeListFirst)}) 第二个最大回撤({calculate_max_drawdown(closeListSecond)})")
    firstPriceList = closeListFirst
    secondPriceList = closeListSecond
    onlyFirst = (firstPriceList[-1] - firstPriceList[returnDiffDay - 1]) / firstPriceList[returnDiffDay - 1]
    onlySecond = (secondPriceList[-1] - secondPriceList[returnDiffDay - 1]) / secondPriceList[returnDiffDay - 1]
    log(
        f'第一个起始日期({firstPriceList.index[returnDiffDay]}) 起始价格({firstPriceList[returnDiffDay]}) 当前日期({firstPriceList.index[-1]}) 当前价格({firstPriceList[-1]})')
    log(
        f'第二个起始日期({secondPriceList.index[returnDiffDay]}) 起始价格({secondPriceList[returnDiffDay]}) 当前日期({secondPriceList.index[-1]}) 当前价格({secondPriceList[-1]})')
    log(f'收益率 仅持有第一个({onlyFirst}) 仅持有第二个({onlySecond})')
    log(f'今天 {firstCode}-{secondCode} {returnDiffDay}日收益差({disList[-1]}%) 收益差总数({len(disList)})\n')
    if needDrawPic:
        drawLastDayPrice(firstCode, secondCode, firstPriceList, secondPriceList, returnDiffDay)

def drawBollLine(firstList, secondList):
    df = pd.DataFrame({'first': firstList, 'second': secondList})
    df = df.fillna(method='bfill')
    df['compare'] = df['first'] / df['second']
    # 计算布林线
    window = 252  # 移动平均窗口
    df['SMA'] = df['compare'].rolling(window=window).mean()  # 中轨线
    df['STD'] = df['compare'].rolling(window=window).std()  # 标准差
    df['Upper'] = df['SMA'] + 2 * df['STD']  # 上轨线
    df['Lower'] = df['SMA'] - 2 * df['STD']  # 下轨线
    plt.figure()
    plt.plot(df['compare'], label='compare', color='blue')
    plt.plot(df['SMA'], label='SMA (Middle Band)', color='orange')
    plt.plot(df['Upper'], label='Upper Band', color='green', linestyle='--')
    plt.plot(df['Lower'], label='Lower Band', color='red', linestyle='--')

    # 填充布林带区域
    #plt.fill_between(df.index, df['Upper'], df['Lower'], color='lightgray', alpha=0.3)

    plt.title('Bollinger Bands')
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.legend()
    plt.show(block=False)
    plt.pause(1)
    bollBackTest(df)


def bollBackTest(df):
    initMoney = 10000
    firstPriceList = df['first']
    firstPriceList = firstPriceList.fillna(method="bfill")
    secondPriceList = df['second']
    secondPriceList = secondPriceList.fillna(method="bfill")
    countFirst = initMoney / 2 / firstPriceList[0]
    countSecond = initMoney / 2 / secondPriceList[0]
    moneyList = []
    moneyList.append(initMoney)
    step = []
    maxMoney = initMoney
    minMoney = initMoney
    for i in range(0, len(firstPriceList)):
        if df['compare'][i] >= df['Upper'][i] and countFirst != 0:
            money = countFirst * firstPriceList[i] * (1 - trade_rate)
            addSecond = money / secondPriceList[i]
            countSecond += addSecond
            countFirst = 0
            step.append(f"{df.index[i]} sell first ")
            hasSwitch = True
            # log("sell first \n")
        if df['compare'][i] <= df['Lower'][i] and countSecond != 0:
            money = countSecond * secondPriceList[i] * (1 - trade_rate)
            addFirst = money / firstPriceList[i]
            countFirst += addFirst
            countSecond = 0
            step.append(f"{df.index[i]} sell second ")
            hasSwitch = True
            # log("sell second\n")
        currentMoney = countFirst * firstPriceList[i] + countSecond * secondPriceList[i]
        moneyList.append(currentMoney)

    money = countFirst * firstPriceList[-1] + countSecond * secondPriceList[-1]
    r = (money - initMoney) / initMoney
    for str in step:
        print(str)
    print(f"boll back test money({money}) r({r}) nianhua({calculate_annualized_return(r,len(df.index))})")



def drawLastDayPrice(firstCode, secondCode, firstList, secondList, diffDay):
    plt.figure()
    firstData = firstList.apply(lambda x: (x - firstList[0]) / firstList[0])
    secondData = secondList.apply(lambda x: (x - secondList[0]) / secondList[0])
    plt.plot(firstData.index, firstData, label=firstCode)
    plt.plot(secondData.index, secondData, label=secondCode)
    plt.grid(True)
    plt.legend()
    plt.savefig(f'image/{firstCode}_{secondCode}_all.jpg', dpi=400)
    plt.show(block=False)
    plt.pause(1)
    plt.figure()
    firstPrice = firstList.tail(diffDay)
    secondPrice = secondList.tail(diffDay)
    firstData = firstPrice.apply(lambda x: (x - firstPrice[0]) / firstPrice[0])
    secondData = secondPrice.apply(lambda x: (x - secondPrice[0]) / secondPrice[0])
    plt.plot(firstPrice.index, firstData, label=firstCode)
    plt.plot(secondPrice.index, secondData, label=secondCode)
    plt.grid(True)
    plt.legend()
    plt.savefig(f'image/{firstCode}_{secondCode}_last.jpg', dpi=400)
    plt.show()


def calculate_annualized_return(total_return, days_held):
    annualized_return = ((1 + total_return / 100) ** (365 / days_held)) - 1
    return round(annualized_return * 100, 2)


def backTest(diffTop, diffBottom, diffList, dateList, firstPriceList, secondPriceList, returnDiffDay):
    initMoney = 10000
    countFirst = initMoney / 2 / firstPriceList[returnDiffDay - 1]
    countSecond = initMoney / 2 / secondPriceList[returnDiffDay - 1]
    moneyList = []
    moneyList.append(initMoney)
    step = []
    maxMoney = initMoney
    minMoney = initMoney
    for i in range(0, len(diffList) - 1):
        if diffList[i] >= diffTop and countFirst != 0:
            money = countFirst * firstPriceList[returnDiffDay + i] * (1 - trade_rate)
            addSecond = money / secondPriceList[returnDiffDay + i]
            countSecond += addSecond
            countFirst = 0
            step.append(f"{dateList[i]} sell first ")
            hasSwitch = True
            # log("sell first \n")
        if diffList[i] <= diffBottom and countSecond != 0:
            money = countSecond * secondPriceList[returnDiffDay + i] * (1 - trade_rate)
            addFirst = money / firstPriceList[returnDiffDay + i]
            countFirst += addFirst
            countSecond = 0
            step.append(f"{dateList[i]} sell second ")
            hasSwitch = True
            # log("sell second\n")
        currentMoney = countFirst * firstPriceList[returnDiffDay + i] + countSecond * secondPriceList[returnDiffDay + i]
        moneyList.append(currentMoney)
        if currentMoney > maxMoney:
            maxMoney = currentMoney
        if currentMoney < minMoney:
            minMoney = currentMoney

    money = countFirst * firstPriceList[-1] + countSecond * secondPriceList[-1]
    r = (money - initMoney) / initMoney
    global gR
    global gTop
    global gBottom
    global gStep
    global gMoney
    if r > gR:
        gR = r
        gTop = diffTop
        gBottom = diffBottom
        gStep = step
        gMoney = moneyList
    # log(f'return r({r})\n')


def drawPic(firstCode, secondCode, dataY, dataX):
    plt.figure()
    time_series = pd.to_datetime(dataX)
    plt.plot(time_series, dataY)
    # 添加标题和标签
    plt.title(f'{firstCode}-{secondCode} 40-day return difference')
    plt.xlabel('time')
    plt.ylabel('diff')
    # 设置 y 轴刻度更密集
    # plt.yticks(range(min(dataY), max(dataY), 1))  # 设置 y 轴刻度间隔为 5
    # 自动旋转 x 轴标签
    plt.xticks(rotation=90)
    # 显示网格线
    plt.grid(True)
    plt.legend()
    plt.savefig(f'image/{firstCode}_{secondCode}.jpg', dpi=400)

    # 显示图形
    plt.show(block=False)
    plt.pause(1)


def drawMoneyPic(firstCode, secondCode, moneyList, dateList):
    plt.figure()
    moneyDf = pd.DataFrame(moneyList, index=dateList)
    ma20List = moneyDf.rolling(window=20).mean()
    time_series = pd.to_datetime(dateList)
    plt.plot(time_series, moneyList)
    plt.plot(time_series, ma20List)
    # 添加标题和标签
    plt.title(f'{firstCode}-{secondCode} money')
    plt.xlabel('time')
    plt.ylabel('money')
    # 设置 y 轴刻度更密集
    # plt.yticks(range(min(dataY), max(dataY), 1))  # 设置 y 轴刻度间隔为 5
    # 自动旋转 x 轴标签
    plt.xticks(rotation=90)
    plt.legend()
    # 显示网格线
    plt.grid(True)
    plt.savefig(f'image/{firstCode}_{secondCode}_money.jpg', dpi=400)

    # 显示图形
    plt.show(block=False)
    plt.pause(1)


def calculate_max_drawdown(prices):
    # 转换为 numpy 数组以便计算
    # prices = np.array(prices)
    # 计算累计最大值
    cumulative_max = np.maximum.accumulate(prices)
    # 计算回撤值
    drawdowns = (prices - cumulative_max) / cumulative_max
    # 找到最大回撤
    max_drawdown = drawdowns.min()
    max_drawdown_end = drawdowns.argmin()
    max_drawdown_start = prices[:max_drawdown_end + 1]
    return max_drawdown


def dayReport():
    # 信息技术 红利 200-50 38% [7~-13]
    log("\n信息技术 红利  38% [7~-13]\n")
    fund('159939', '512890')


def log(msg):
    logger.info(msg)


if __name__ == '__main__':
    logger.info("starting...")
    # dayReport()
    # 信息技术 红利低波 年化45% 【13，-23】适合20日、30日
    fund('159939', '512890')
    # 信息技术 酒 年化33% 【39，-30】
    # fund('159939', '512690')
    # 科创50 红利低波 年化60% 【9，-24】适合30日，40日 981
    #fund('588000', '512890', needDrawPic=True)
    # 信息技术 食品饮料 年化21% 【4，-9】
    # fund('159939', '515170')
    # 信息技术 红利 年化33% 【4，-25】2389
    #fund('159939', '510880')
    # 科创50 红利 年化62% 【8，-25】 981
    # fund('588000', '510880')
    # 沪深300 中证1000 年化12% 【15，-1】 1956
    # fund('510300', '512100')
    # 沪深300 创业板 年化12% 【15，-1】 1956
    # fund('510300', '159915')
